// udp_rpc_client.go
//
// 功能：模拟并发客户端调用 UDP JSON-RPC 服务。
// 流程：
//   1. 使用 UDP 连接服务器。
//   2. 启动多个 goroutine，每个 goroutine 发送若干请求，随机生成加数。
//   3. 为每个请求生成唯一 RequestID，并等待服务器响应。
//   4. 使用 WaitGroup 确保所有请求完成后退出。
// 关键点：
//   - UDP 无连接，需要在客户端维护 requestID -> 响应的匹配关系。
//   - 通过 channel 和 map 结合来追踪每个请求的结果。
//   - 加入超时机制，防止丢包导致 goroutine 永久阻塞。
//   - 所有日志均标注 requestID，便于调试。

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"math"
	"math/rand"
	"net"
	"sync"
	"time"
)

const (
	serverAddr       = "127.0.0.1:9000" // 服务端地址
	clientTimeout    = 2 * time.Second  // 单个请求等待响应的超时时间
	concurrentClient = 5                // 并发 goroutine 数量
	requestPerClient = 20               // 每个 goroutine 发送的请求数
)

// AddRequest/Response 与服务器结构体保持一致。
type AddRequest struct {
	RequestID string  `json:"request_id"`
	A         float64 `json:"a"`
	B         float64 `json:"b"`
}

type AddResponse struct {
	RequestID string  `json:"request_id"`
	Sum       float64 `json:"sum"`
	Error     string  `json:"error,omitempty"`
}

// responseRegistry 管理 requestID 到结果通道的映射，确保并发安全。
type responseRegistry struct {
	mu    sync.Mutex
	cache map[string]chan AddResponse
}

func newRegistry() *responseRegistry {
	return &responseRegistry{
		cache: make(map[string]chan AddResponse),
	}
}

func (r *responseRegistry) register(id string) chan AddResponse {
	r.mu.Lock()
	defer r.mu.Unlock()
	ch := make(chan AddResponse, 1)
	r.cache[id] = ch
	return ch
}

func (r *responseRegistry) notify(resp AddResponse) {
	r.mu.Lock()
	ch, ok := r.cache[resp.RequestID]
	if ok {
		delete(r.cache, resp.RequestID)
	}
	r.mu.Unlock()

	if ok {
		ch <- resp
		close(ch)
	}
}

func main() {
	rand.Seed(time.Now().UnixNano())

	serverUDPAddr, err := net.ResolveUDPAddr("udp", serverAddr)
	if err != nil {
		log.Fatalf("解析服务器地址失败: %v", err)
	}

	// 使用 DialUDP 便于读取响应（自动记住远端地址）
	conn, err := net.DialUDP("udp", nil, serverUDPAddr)
	if err != nil {
		log.Fatalf("建立 UDP 连接失败: %v", err)
	}
	defer conn.Close()

	log.Printf("[client] 已连接到 %s", serverAddr)

	registry := newRegistry()
	var wg sync.WaitGroup

	// 负责接收服务器响应的专用 goroutine。
	go receiveLoop(conn, registry)

	start := time.Now()
	for i := 0; i < concurrentClient; i++ {
		clientID := i
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < requestPerClient; j++ {
				reqID := fmt.Sprintf("client%d-%d-%d", clientID, j, time.Now().UnixNano())
				a := rand.Float64() * 100
				b := rand.Float64() * 100
				if err := sendRequest(conn, registry, reqID, a, b); err != nil {
					log.Printf("[client] 请求失败 id=%s err=%v", reqID, err)
				}
				time.Sleep(20 * time.Millisecond)
			}
		}()
	}

	wg.Wait()
	elapsed := time.Since(start)
	log.Printf("[client] 全部请求发送完成，总耗时 %s", elapsed)
	time.Sleep(500 * time.Millisecond) // 等待接收协程输出完毕
}

// sendRequest 负责发送请求并等待响应。
func sendRequest(conn *net.UDPConn, registry *responseRegistry, id string, a, b float64) error {
	req := AddRequest{
		RequestID: id,
		A:         a,
		B:         b,
	}
	payload, err := json.Marshal(req)
	if err != nil {
		return fmt.Errorf("序列化请求失败: %w", err)
	}

	respCh := registry.register(id)

	if _, err := conn.Write(payload); err != nil {
		return fmt.Errorf("发送请求失败: %w", err)
	}

	select {
	case resp := <-respCh:
		if resp.Error != "" {
			return fmt.Errorf("服务器返回错误: %s", resp.Error)
		}
		if !floatEquals(resp.Sum, a+b) {
			return fmt.Errorf("结果不一致，期望=%v 实际=%v", a+b, resp.Sum)
		}
		log.Printf("[client] 请求成功 id=%s a=%v b=%v sum=%v", id, a, b, resp.Sum)
	case <-time.After(clientTimeout):
		return fmt.Errorf("请求超时 id=%s", id)
	}
	return nil
}

// receiveLoop 持续读取服务器响应，根据 requestID 分发。
func receiveLoop(conn *net.UDPConn, registry *responseRegistry) {
	buf := make([]byte, 4096)
	for {
		n, err := conn.Read(buf)
		if err != nil {
			log.Printf("[client] 读取响应失败: %v", err)
			return
		}
		var resp AddResponse
		if err := json.Unmarshal(buf[:n], &resp); err != nil {
			log.Printf("[client] 响应 JSON 解析失败: %v 原始=%q", err, string(buf[:n]))
			continue
		}
		registry.notify(resp)
	}
}

// floatEquals 对浮点数求和做近似比较。
func floatEquals(a, b float64) bool {
	const epsilon = 1e-6
	return math.Abs(a-b) < epsilon
}
